#include <stdio.h>
#include <mpi.h>

#define F(x) x*x

void get_input(int my_rank, double* a_p,  double* b_p, int* n_p);
double trapeziumRule(double a, double b, int n , double h);

int main(int argc, char **argv)
{
    int  comm_sz, my_rank;
    int n, local_n;
    double a, b, h, local_a, local_b, local_integral;

    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &comm_sz);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);

    get_input(my_rank, comm_sz, &a, &b, &n);

    h = (b-a)/(double)n;
    local_n = n/comm_sz;

    int mod = n % comm_sz;
    if (mod != 0) {
        if (n >  comm_sz) {
            if (my_rank < mod) {
                /* Distribute the mod to the first ones */
                local_n ++;
                local_a = a + local_n*my_rank*h;
                local_b = local_a + local_n*h;
            } else { /* (my_rank > mod) */
                local_a = a + (local_n +1)*my_rank*h - (my_rank - mod)*h;
                local_b = local_a + local_n*h;
            }
            local_integral = trapeziumRule(local_a, local_b, local_n, h);

        } else { /* (n < comm_sz) and (n/comm_sz == 0) */
            if (my_rank < mod) {
                /* Allow the first ones to execute */
                local_n ++;
                local_a = a + local_n*my_rank*h;
                local_b = local_a + local_n*h;
                local_integral = trapeziumRule(local_a, local_b, local_n, h);
            } else { /* (my_rank > mod) */
                local_integral = 0;
            }
        }

    } else { /* (mod == 0)*/
        local_a = a + my_rank*local_n*h;
        local_b = local_a + local_n*h;
        local_integral = trapeziumRule(local_a, local_b, local_n, h);
    }

    double total_integral;
    MPI_Reduce(&local_integral, &total_integral, 1, MPI_DOUBLE, MPI_SUM, 0,
               MPI_COMM_WORLD);

    if (my_rank == 0) {
        printf("The integral is: %f\n", total_integral);
    }
    MPI_Finalize();
    return 0;
}

void get_input(int my_rank, double* a_p,  double* b_p, int* n_p)
{
    if (my_rank == 0) {
        printf("Enter a, b, and n\n");
        scanf("%lf %lf %d", a_p, b_p, n_p);
    }
    MPI_Bcast(a_p, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(b_p, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(n_p, 1, MPI_INT, 0, MPI_COMM_WORLD);
}

double trapeziumRule(double a, double b, int n, double h)
{
    double x_i, approx;
    approx = (F(a) + F(b))/2.0;

    int i;
    for (i=1; i<=n-1; i++) {
        x_i = a + i*h;
        approx += F(x_i);
    }
    approx = h*approx;
    return approx;
}
